import random
from os.path import join, abspath
from typing import List
from backend.experiment.framework.shots_prompt.implement import \
    CommonPerplexityChoiceMakingSample
import json
from backend.experiment.clue.c3.scheme0.implement import C3Implement as S0


SAMPLE_LIMIT = 50


class C3Implement(S0):
    @classmethod
    def _read_train_records(cls) -> list:
        with open(join(cls.data_dir(), 'm-train.json'), 'r', encoding='utf8') \
                as file:
            records = json.loads(file.read())
        with open(join(cls.data_dir(), 'd-train.json'), 'r', encoding='utf8') \
                as file:
            records.extend(json.loads(file.read()))
        random.shuffle(records)
        return records[:SAMPLE_LIMIT]

    @classmethod
    def _make_test_sample(cls, record, made_sample_num: int) -> \
            List[CommonPerplexityChoiceMakingSample]:
        res = []
        text = '\n'.join(record[0])
        questions = record[1]
        for each in questions:
            question = each['question']
            choice = each['choice']
            prompt = cls.__make_single_prompt(
                text, question, cls.__mask_token)
            res.append(
                CommonPerplexityChoiceMakingSample(
                    prompt=prompt,
                    target_mask_token=cls.__mask_token,
                    sample_id=made_sample_num,
                    candidates=choice
                )
            )
            made_sample_num += 1
        return res

    @classmethod
    def cache_dir(cls) -> str:
        return join(cls.work_dir(), 'cache')

    @classmethod
    def work_dir(cls) -> str:
        return abspath(join(__file__, '..'))
